package utils

import (
	"crypto"
	"crypto/subtle"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"syscall"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

const (
	RSABlockType          = "RSA PRIVATE KEY"
	ECDSABlockType        = "EC PRIVATE KEY"
	SM2EncryptedBlockType = "ENCRYPTED SM2 PRIVATE KEY"
	SM2BlockType          = "SM2 PRIVATE KEY"
)

// KeysEqual is constant time compare of the keys to avoid timing attacks.
func KeysEqual(ak, bk ssh.PublicKey) bool {

	//avoid panic if one of the keys is nil, return false instead
	if ak == nil || bk == nil {
		return false
	}

	a := ak.Marshal()
	b := bk.Marshal()
	return (len(a) == len(b) && subtle.ConstantTimeCompare(a, b) == 1)
}

type publicKey struct {
	comments string
	marshal  string
	key      ssh.PublicKey
}

func LoadSSHPubKeyBytes(pubBytes []byte) (map[string][]ssh.PublicKey, error) {
	authorizedKeysMap := map[string]*publicKey{}
	for len(pubBytes) > 0 {
		pubKey, comments, _, rest, err := ssh.ParseAuthorizedKey(pubBytes)
		if err != nil {
			return nil, err
		}
		pub := &publicKey{
			comments: comments,
			marshal:  string(pubKey.Marshal()),
			key:      pubKey,
		}
		authorizedKeysMap[string(pubKey.Marshal())] = pub
		pubBytes = rest
	}

	pubs := make(map[string][]ssh.PublicKey)
	for _, v := range authorizedKeysMap {
		var p []ssh.PublicKey
		if len(v.comments) > 0 {
			if val, ok := pubs[v.comments]; ok {
				p = append(val, v.key)
			} else {
				p = append(p, v.key)
			}
			pubs[v.comments] = p
		} else {
			// default is user cert
			if val, ok := pubs["user"]; ok {
				p = append(val, v.key)
			} else {
				p = append(p, v.key)
			}
			pubs["user"] = p
		}
	}
	return pubs, nil
}

func LoadSSHPubKey(keyPath string) (map[string][]ssh.PublicKey, error) {
	authorizedKeysBytes, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, err
	}
	return LoadSSHPubKeyBytes(authorizedKeysBytes)
}

func LoadkeyBytes(keyBytes []byte, password []byte) (crypto.Signer, error) {
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		return nil, fmt.Errorf("decode private key error")
	}
	var der []byte
	if x509.IsEncryptedPEMBlock(block) && password == nil {
		fmt.Printf("Enter passphrase for (empty for no passphrase):\n")
		password, err := terminal.ReadPassword(int(syscall.Stdin))
		if err != nil {
			return nil, err
		}
		der, err = x509.DecryptPEMBlock(block, password)
	} else {
		der = block.Bytes
	}

	if block.Type == ECDSABlockType {
		return x509.ParseECPrivateKey(der)
	} else if block.Type == RSABlockType {
		return x509.ParseECPrivateKey(der)
	}
	return nil, fmt.Errorf("unsupported block type %s", block.Type)
}

func ListenUnix(name string) (net.Listener, error) {
	socketPath, err := GetUnixSockPath(name)
	if err != nil {
		return nil, err
	}
	l, err := net.Listen("unix", socketPath)
	return l, err
}

func GetUnixSockPath(name string) (string, error) {
	socketDir, err := ioutil.TempDir(os.TempDir(), "ragdoll-")
	if err != nil {
		return "", err
	}
	return filepath.Join(socketDir, fmt.Sprintf("ragdoll-%v.socket", name)), nil
}
